import torch


def mlm_loss(logits, mask_positions, sub_mask_labels, cross_entropy_criterion, device):
    # logits: 模型的预测输出，形状为 (batch, seq_len, vocab_size)。表示每个位置上所有词汇的概率分布。
    # mask_positions: 每个样本中 [MASK] 的位置，形状为 (batch, mask_label_num)。
    # sub_mask_labels: 每个 [MASK] 的真实标签，可能包含多个子标签。
    """
    计算指定位置的mask token的output与label之间的cross entropy loss。
    对于有多个子标签的父标签，将预测标签值复制len(父标签)份【repeat】,再进行维度转换与真实子标签进行损失计算求平均
    Args:
        logits (torch.tensor): 模型预测输出 -> (batch, seq_len, vocab_size)
        mask_positions (torch.tensor): mask token的位置  -> (batch, mask_label_num)
        sub_mask_labels (list): mask token的sub label, 由于每个label的sub_label数目不同，所以这里是个变长的list,
                            e.g. -> [
                                        [[2398, 3352]],
                                        [[2398, 3352], [3819, 3861]]
                                    ]
        cross_entropy_criterion (CrossEntropyLoss): CE Loss计算器
        device (str): cpu还是gpu
    Returns:
        torch.tensor: CE Loss(一个批次的平均损失)
    """
    # 获取形状
    batch_size, seq_len, vocab_size = logits.size()
    loss = 0
    # 遍历每个样本计算损失
    for values in zip(logits, sub_mask_labels, mask_positions):
        # print('*' * 50)

        # s_logits：当前样本的预测值，形状为 (seq_len, vocab_size)。s_labels：当前样本的子标签。s_position：当前样本的掩码位置。
        s_logits, s_labels, s_position = values

        # 获取mask所在位置的预测值
        s_mask_logits = s_logits[s_position]

        # 如果一个 [MASK] 有多个子标签，需要将预测值复制多份，以便计算每个子标签的损失。
        s_mask_logits = s_mask_logits.repeat(len(s_labels), 1, 1)
        # print(s_mask_logits.shape)
        # 预测结果：形状调整
        s_mask_logits = s_mask_logits.reshape(-1, vocab_size)
        # print(s_mask_logits.shape)
        # 真实值
        s_labels = torch.LongTensor(s_labels).to(device)
        # print('s_labels.shape',s_labels.shape)

        s_labels = s_labels.reshape(-1, 1).squeeze()
        # print('s_labels.shape', s_labels.shape)

        # 这段代码用于处理当s_labels在某些情况下维度缺失的问题。如果s_labels的维度为0（即为空或标量），
        # 则通过unsqueeze(dim=0)为其增加一个维度，确保后续计算时形状一致，避免报错。
        if not s_labels.size():  # 处理单token维度下维度缺失的问题
            s_labels = s_labels.unsqueeze(dim=0)
        # 如果多个子标签，算平均损失
        s_loss = cross_entropy_criterion(s_mask_logits, s_labels) / len(s_labels)
        loss += s_loss
    # 求一个批次的平均损失
    loss = loss / batch_size
    return loss


def convert_logits_to_ids(logits: torch.tensor, mask_positions: torch.tensor):
    """
    输入模型预测的词表概率分布（LMModel的logits），将mask_position位置的token logits转换为token的id。
    Args:
        logits (torch.tensor): model output -> (batch, seq_len, vocab_size)
        mask_positions (torch.tensor): mask token的位置 -> (batch, mask_label_num)
    Returns:
        torch.LongTensor: 对应mask position上最大概率的推理token(子标签id) -> (batch, mask_label_num)
    """
    # 获取形状
    batch_size, seq_len, vocab_size = logits.size()
    # print('batch_size:{} seq_len:{} vocab_size:{}'.format(batch_size, seq_len, vocab_size))
    batch_size, label_len = mask_positions.size()
    # print('batch_size:{} label_len:{}'.format(batch_size, label_len))

    # 改变logits的形状
    logits = logits.reshape(-1, vocab_size)
    # print('logits.shape:', logits.shape)

    # 获取铺平之后的 mask_position
    # [[5, 6], [7, 8]] -> [5, 6, 27, 28]
    mask_positions_reshape = []
    # print('mask_positions:', mask_positions.detach().numpy().tolist())
    for batch_id, mask_position in enumerate(mask_positions.detach().numpy().tolist()):
        # print('batch_id:', batch_id)
        # print('mask_position:', mask_position)
        for pos in mask_position:
            # print('pos:', pos)
            mask_positions_reshape.append(batch_id * seq_len + pos)
    # print('mask_positions_reshape:', mask_positions_reshape) # [5, 6, 27, 28]

    # 通过切片获取mask的logits
    mask_logits = logits[mask_positions_reshape]
    # print('mask_logits.shape:', mask_logits.shape)
    # 获取最大概率的token id
    ids = torch.argmax(mask_logits, dim=1)
    # print('ids:', ids)
    # print('ids.shape:', ids.shape)
    # 改变形状，返回结果
    ids = ids.reshape(-1, label_len)
    # print('ids:', ids)
    # print('ids.shape:', ids.shape)
    return ids


if __name__ == '__main__':
    logits = torch.randn(2, 20, 1000)
    print('logits.shape:', logits.shape)  # torch.Size([2, 20, 1000])
    mask_positions = torch.LongTensor([[5, 6], [7, 8]])  # torch.Size([2, 2])
    print('mask_positions.shape:', mask_positions.shape)
    sub_mask_labels = [[[50, 60]], [[72, 83], [34, 65]]]

    CE_loss = torch.nn.CrossEntropyLoss()
    device = 'cpu'

    loss = mlm_loss(logits, mask_positions, sub_mask_labels, CE_loss, device)
    print(loss)

    # 从logits找出mask_position位置的token
    ids = convert_logits_to_ids(logits, mask_positions)
    print(ids)
